#include "tree.h"
#include <time.h>
#include <assert.h>
#include <string.h>
#include <stdlib.h>
#include <queue>

typedef unsigned char byte;

int NBytes(int nbits) { return (nbits + 7) / 8; }

byte *GenBytes(int n) {
  byte *bytes = new byte[n];
  for (int i = 0; i < n; ++i) {
    bytes[i] = byte((rand() % 256));
  }
  return bytes;
}

void CharFrequency(const byte bytes[], int n, int frequency[]) {
  printf("count frequency ...\n");
  memset(frequency, 0, sizeof(int) * 256);
  for (int i = 0; i < n; ++i) ++frequency[bytes[i]];
}

void CharFrequency(FILE *fp, int frequency[]) {
}

struct FrequencyNode {
  FrequencyNode(int freq, Node *node) : freq(freq), node(node) { }
  FrequencyNode() : freq(0), node(NULL) { }
  bool operator<(const FrequencyNode &n) const {
    return this->freq > n.freq;
  }
  int freq;
  Node *node;
};

typedef std::priority_queue<FrequencyNode> Queue;

Node *BuildHuffmanTree(int frequency[]) {
  printf("build Huffman tree ...\n");

  Queue q;
  for (int i = 0; i < 256; ++i) {
    if(frequency[i] > 0) {
      q.push(FrequencyNode(frequency[i], new Node(i)));
    }
  }
  if (q.empty()) return NULL;

  while (q.size() > 1) {
    FrequencyNode min1 = q.top();
    q.pop();
    FrequencyNode min2 = q.top();
    q.pop();
    //printf("%d, %d, %d\n", q.size(), min1.freq, min2.freq);

    Node *node = new Node();
    node->left = min1.node;
    node->right = min2.node;
    q.push(FrequencyNode(min1.freq + min2.freq, node));
  }
  Node *root = q.top().node;
  q.pop();

  return root;
}

struct HuffmanCode {
  HuffmanCode() : val(0), n(0) { }
  HuffmanCode(int val, byte n) : val(val), n(n) { }

  // support at most 32 bits code
  int val;
  byte n;
};

void RetrieveCode(const Node *node, int code, byte n, HuffmanCode codeset[]) {
  if (node->IsLeaf()) {
    codeset[node->key] = HuffmanCode(code, n);
  }
  else {
    assert(node->left != NULL && node->right != NULL);

    RetrieveCode(node->left, code | 0, n + 1, codeset);
    RetrieveCode(node->right, code | (1 << n) , n + 1, codeset);
  }
}

void RetrieveCode(const Node *root, HuffmanCode codeset[]) {
  printf("retrieve Huffman code ...\n");

  if (root == NULL) return;
  RetrieveCode(root->left, 0, 1, codeset);
  RetrieveCode(root->right, 1, 1, codeset);
}

void GenHuffmanCode(int frequency[], HuffmanCode code[]) {
  Node *root = BuildHuffmanTree(frequency);
  RetrieveCode(root, code);
  DeleteTree(&root);
}

int StoreCode(byte bytes[], int bit, const HuffmanCode &code) {
  int i = bit / 8;
  int j = bit % 8;
    
  int remained_bits = code.val;
  // store the first (8 - j) bits to bytes[i] at j
  bytes[i] |= remained_bits << j; 
  
  remained_bits >>= 8 - j;
  int n_remained_bits = code.n - (8 - j);
  while (n_remained_bits >= 8) {
    bytes[++i] = remained_bits & 0xff;
    remained_bits >>= 8;
    n_remained_bits -= 8;
  }
  if (n_remained_bits > 0) bytes[i + 1] |= remained_bits;
  
  return bit + code.n;
}

int EncodedSize(const HuffmanCode codeset[], const int freq[]) {
  int nbits = 0;
  for (int i = 0; i < 256; ++i) {
    nbits += freq[i] * codeset[i].n;
  }
  return NBytes(nbits);
}

byte *Encode(const HuffmanCode codeset[], const byte bytes[], int n, 
             int nbytes) {  
  byte *encoded = new byte[nbytes];
  memset(encoded, 0, sizeof(byte) * nbytes);

  int bit = 0;
  for (int i = 0; i < n; ++i) {   
    bit = StoreCode(encoded, bit, codeset[bytes[i]]);
  }  
  
  return encoded;
}

byte Restore(const Node *root, const byte *encoded, int *bit) {
  const Node *p = root;
  while (!p->IsLeaf()) {
    int i = (*bit) / 8;
    int j = (*bit) % 8;
    if ((encoded[i] & (1 << j)) == 0) {
      p = p->left;
    }
    else {
      p = p->right;
    }
    ++(*bit);
  }
  return p->key;
}

byte *Decode(const Node *root, const byte *encoded, int n) {
  if (root == NULL) return NULL;

  byte *restored = new byte[n];
  for (int i = 0, bit = 0; i < n; ++i) {
    restored[i] = Restore(root, encoded, &bit);
  }
  return restored;
}

void Print(const HuffmanCode codeset[], const int freq[]) {
  for (int i = 0; i < 256; ++i) {
    if (freq[i] > 0) {
      printf("%d: %d, %d, %x\n", i, freq[i], codeset[i].n, codeset[i].val);
    }
  }
}

void Print(const byte bytes[], int n) {
  for (int i = 0; i < n; ++i) printf("%c", bytes[i]);
  printf("\n");
}

void Check(const byte bytes1[], const byte bytes2[], int n) {
  for (int i = 0; i < n; ++i) {
    if (bytes1[i] != bytes2[i]) {
      printf("%d unmatched: %d, %d\n", i, bytes1[i], bytes2[i]);
      return;
    }
  }
  printf("ok\n");
}

void Test(const byte *bytes, int n) {
  //Print(bytes, n);

  int freq[256];
  CharFrequency(bytes, n, freq);

  Node *root = BuildHuffmanTree(freq);

  HuffmanCode codeset[256];
  RetrieveCode(root, codeset);
  //Print(codeset, freq);

  int nbytes = EncodedSize(codeset, freq);
  byte *encoded = Encode(codeset, bytes, n, nbytes);

  byte *restored = Decode(root, encoded, n);
  Print(restored, n);

  Check(bytes, restored, n);

  printf("original size: %d (bytes) \ncompressed size: %d (bytes)\n",
         n, nbytes);

  DeleteTree(&root);
  delete[] restored;
  delete[] encoded;
}

void Test(int n) {
  printf("test, n = %d\n", n);

  byte *bytes = GenBytes(n);
  Test(bytes, n);
  
  delete[] bytes;
}

void Test(const char *filename) {
  FILE *fp = fopen(filename, "rb");
  fseek(fp, 0, SEEK_END);
  int nbytes = ftell(fp);
  byte *bytes = new byte[nbytes];
  fseek(fp, 0, SEEK_SET);
  fread(bytes, 1, nbytes, fp);
  fclose(fp);

  Test(bytes, nbytes);

  delete[] bytes;  
}

int main(int argc, const char *argv[]) {
  srand(time(NULL));
  if (argc < 2) {
    Test(1000);
  }
  else {
    Test(argv[1]);
  }
  return 0;
}

